# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging

from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.subscription import SubscriptionRequest, FeedSubscriptionRequest
from coin.feed.fastfeed import feed_pb2
from coin.feed.fastfeed.proto_feed_handler import create_proto_feed_handler
from coin.feed.fastfeed.feed_cache import FeedCacheKey, FeedCacheWriter
from coin.feed.fastfeed.util import (
    get_products_from_sub_req,
)


def subscribe_all_feed_types(sub_req, products, proto_feed_handler):
  if sub_req == SubscriptionRequest('Futures', 'Bitmex'):
    for product in products:
      proto_feed_handler.subscribe_book(product, None)
      proto_feed_handler.subscribe_trade(product, None)
      proto_feed_handler.subscribe_instrument(product, None)

  elif (sub_req == SubscriptionRequest('Futures', 'Okex', 'v3')
        or sub_req == SubscriptionRequest('Futures', 'Okex', 'swap')):
    for product in products:
      proto_feed_handler.subscribe_book(product, None)
      proto_feed_handler.subscribe_trade(product, None)
      proto_feed_handler.subscribe_instrument(product, None)

  else:
    for product in products:
      proto_feed_handler.subscribe_book(product, None)
      proto_feed_handler.subscribe_trade(product, None)


class FeedCacheGenerator:
  def __init__(self,
               machine,
               date,
               sub_req,
               worker_id,
               feed_cache_dir,
               products=None,
               quiet=False,
               book_builder_opts=None):
    self.feed_cache_key = FeedCacheKey(machine=machine,
                                       date=date,
                                       sub_req=sub_req,
                                       worker_id=worker_id)
    self.feed_cache_dir = feed_cache_dir

    self.products = products if products is not None else get_products_from_sub_req(sub_req, date)
    self.feed_sub_req = FeedSubscriptionRequest.create(self.products, sub_req)

    self._proto_feed_handler = None
    self._feed_cache_writer = None
    self._last_feed_ts = 0
    self._last_log_ts = 0
    self._quiet = quiet
    self._book_builder_opts = book_builder_opts

  def run(self):
    assert self._proto_feed_handler is None
    from_ts = to_datetime(self.feed_cache_key.date)
    to_ts = from_ts + datetime.timedelta(hours=24, minutes=5)

    if self.feed_cache_key.worker_id == '2':
      from_ts -= datetime.timedelta(minutes=30)
    logging.info('Conversion: %s - %s', from_ts, to_ts)

    self._feed_cache_writer = FeedCacheWriter(self.feed_cache_key,
                                              feed_cache_dir=self.feed_cache_dir)

    try:
      run_from_archive(self.feed_sub_req,
                       self.on_book_reset,
                       from_ts,
                       to_ts,
                       machine=self.feed_cache_key.machine,
                       worker_id=self.feed_cache_key.worker_id,
                       use_feed_arbitration=None,
                       book_builder_opts=self._book_builder_opts)
    except Exception:
      logging.exception('Exception occured at %s', to_datetime(self._last_feed_ts))
      raise

    self._feed_cache_writer.close()
    self._feed_cache_writer = None

  def on_book_reset(self, sub_req, underlying_feed_handler):
    assert self._proto_feed_handler is None
    self._proto_feed_handler = create_proto_feed_handler(sub_req,
                                                         underlying_feed_handler,
                                                         no_parse=True)
    self._proto_feed_handler.subscribe_proto_feed(self.on_proto_feed)
    subscribe_all_feed_types(sub_req, self.products, self._proto_feed_handler)

  def on_proto_feed(self, product, msgs):
    if not msgs:
      return

    ts = msgs[0].timestamp
    self._last_feed_ts = ts
    if not self._quiet and ts - self._last_log_ts >= 30 * 60 * 10**9:
      logging.info(to_datetime(ts))
      self._last_log_ts = ts

    for msg in msgs:
      if msg.feed_type != feed_pb2.FEED_TYPE_METADATA and msg.sort_key == 0:
        logging.debug('sort_key is zero: \n%s', str(msg))

    self._feed_cache_writer.write_proto_msg(product, msgs)


def main(argv):
  FLAGS = flags.FLAGS
  sub_req = SubscriptionRequest.from_str(FLAGS.sub_req)
  date = datetime.datetime.strptime(FLAGS.date, '%Y%m%d').date()

  logging.info('---\nDate: %s\nMachine: %s\nSubReq: %s\nWorker ID: %s',
               str(date),
               FLAGS.machine,
               FLAGS.sub_req,
               FLAGS.worker_id)

  all_products = get_products_from_sub_req(sub_req, date, FLAGS.group)
  products = []
  for p in all_products:
    if ((not FLAGS.base_filter or p.base.currency in FLAGS.base_filter)
        and (not FLAGS.quote_filter or p.quote.currency in FLAGS.quote_filter)):
      products.append(p)
    elif (FLAGS.symbol_filter and p.symbol in FLAGS.symbol_filter):
      products.append(p)

  book_builder_opts = {}
  if FLAGS.crop_book_snapshot_depth is not None:
    book_builder_opts['crop_book_snapshot_depth'] = FLAGS.crop_book_snapshot_depth
  feed_cache_gen = FeedCacheGenerator(machine=FLAGS.machine,
                                      date=date,
                                      sub_req=sub_req,
                                      worker_id=FLAGS.worker_id,
                                      feed_cache_dir=FLAGS.feed_cache_dir,
                                      products=products,
                                      quiet=FLAGS.quiet,
                                      book_builder_opts=book_builder_opts)
  feed_cache_gen.run()


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string('machine', 'feed-01.ap-northeast-1.aws', 'machine')
  flags.DEFINE_string('date', None, 'Date')
  flags.DEFINE_string(
      'sub_req',
      'Futures.Bitmex',
      'Subscription Request String. <market_type>.<exchange>.<api_version> or '
      '<market_type>.<exchange>')
  flags.DEFINE_string('worker_id', '1', 'Worker ID')
  flags.DEFINE_string('feed_cache_dir', '~/data/feed_cache', 'Feed cache directory')
  flags.DEFINE_bool('quiet', False, 'Quite')
  flags.DEFINE_integer('crop_book_snapshot_depth', None, 'Crop book snapshot')

  flags.DEFINE_list('base_filter',
                    [
                        'BTC',
                        'ETH',
                        'LTC',
                        'ETC',
                        'BCHABC',
                        'XRP',
                        'EOS',
                        'BNB',
                        'ADA',
                        'BCHSV',
                        'TRX',
                        'LINK',
                        'XLM',
                        'STEEM',
                        'TMTG',
                        'FCT',
                        'BHP',
                        'THETA',
                        'RNT',
                        'FNB',
                        'INS',
                        'META',
                        'EL',
                        'TON',
                        'SXP',
                    ],
                    'Quote currency filter')
  flags.DEFINE_list(
      'quote_filter',
      ['USD', 'KRW', 'JPY', 'USDT', 'USDC', 'TUSD', 'PAX', 'USDK', 'BTC', 'ETH', 'EUR'],
      'Base curency filter')
  flags.DEFINE_list('symbol_filter', ['BTC-HUSD'], 'symbol filter')
  flags.DEFINE_string('group', 'all', 'Empty group means all groups.')
  app.run(main)
